package com.shijie.sjoj.service.box;

import com.shijie.sjoj.model.component.ParamsTypeFormat;
import com.shijie.sjoj.model.constant.ExceptionConstant;
import com.shijie.sjoj.model.exception.JudgeException;
import com.shijie.sjoj.model.pojo.base.CodeType;
import com.shijie.sjoj.model.pojo.base.JudgedInfo;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.asm.ClassReader;
import org.springframework.asm.ClassVisitor;
import org.springframework.asm.MethodVisitor;
import org.springframework.asm.Opcodes;
import org.springframework.stereotype.Component;

import javax.tools.*;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.management.ManagementFactory;
import java.lang.management.MemoryMXBean;
import java.lang.management.MemoryUsage;
import java.lang.reflect.Method;
import java.lang.reflect.ReflectPermission;
import java.net.URI;
import java.net.URL;
import java.net.URLClassLoader;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.*;
import java.util.*;

/**
 * Java语言判题机
 *
 * @author ZSJ
 */
@Slf4j
@Component
@NoArgsConstructor
@JudgeCodeType(value = CodeType.Java)
public class JudgeWorkerJava implements JudgeWorker {

    @Override
    public JudgedInfo working(String code, CodeType codeType,
                              Class<?>[] argsType, List<List<?>> args,
                              Class<?> outcomeType, List<List<?>> outcome) {
        // 获取代码的类名和方法名
        String[] nameInfo = this.getClassNameAndMethodName(code);
        String className = nameInfo[0];
        String methodName = nameInfo[1];
        // 设置编译结果的输出目录
        String outDir = JudgeWorker.OutDirPrefix + codeType.getValue();
        // 删除之前可能存在的编译文件
        try {
            Files.deleteIfExists(Paths.get(outDir, className + ".class"));
        } catch (IOException e) {
            log.warn(e.toString());
        }
        // 编译代码
        this.compiler(code, className, outDir);
        // 加载类（字节码）文件并输出最终结果
        return this.classLoadExecute(outDir, className, methodName, argsType, args, outcomeType, outcome);
    }

    /**
     * 获取代码的类名（className）和方法名（methodName）
     *
     * @param code 代码
     * @return String[] index-0：类名，index-1：方法名
     */
    private String[] getClassNameAndMethodName(String code) {
        // 规定：1.类与导入的包空一行，2.类与方法都为public
        // 默认的类名与方法名
        String className = "SolveClass";
        String methodName = "solving";
        // 拆分代码
        String[] split = code.split(" ");
        // 最简形式类[必须写修饰符]：public class Clazz{public void run(){}}
        if (split.length < 5) {
            throw JudgeException.throwError(ExceptionConstant.CODE_MISTAKE);
        }
        // 开始获取类名与方法名
        for (int i = 0; i < split.length; i++) {
            // class关键字后第一个不为空的就是类名所在串
            if (split[i].contains("class")) {
                // 移至class后的串并排除可能的空串
                do i++;
                while (split[i].isBlank());
                // 此时可能的几种结果:【 MyClass, MyClass{, MyClass{public 】，
                // 进一步处理，第一个\将(第二个)\转义，第二个\转义{
                // 获取类名，trim() 去除可能因文本块造成的空白
                className = split[i].split("\\{")[0].trim();

                // 找到包含public关键字的位置，方法名在其第两个非空串处
                while (!split[i].contains("public")) i++;
                // 到达public后的第一个非空串后的串并排除可能的[第一列]空串
                do i++;
                while (split[i].isBlank());
                // 到达public后的第二个非空串后的串并排除可能的[第二列]空串
                do i++;
                while (split[i].isBlank());
                // 此时可能的几种结果:【 name, name(, name(Object 】，
                // 进一步处理，第一个\将(第二个)\转义，第二个\转义(
                // 获取方法名，trim() 去除可能因文本块造成的空白
                methodName = split[i].split("\\(")[0].trim();
                break;
            }
        }
        // 返回类名和方法名
        return new String[]{className, methodName};
    }

    /**
     * 编译代码
     *
     * @param code      代码
     * @param className 代码的类名
     * @param outDir    代码的方法名
     */
    private void compiler(String code, String className, String outDir) {
        // 将源码code转为JavaFileObject对象准备编译
        // 或者通过文件获取javaFileObjects fileManager.getJavaFileObjectsFromPaths(List.of())
        JavaFileObject javaFileObject = new SimpleJavaFileObject(
                URI.create("string://code-java/" + className + ".java"), // 虚拟 URI
                JavaFileObject.Kind.SOURCE // 指定文件类型
        ) {
            //重写方法，获取待转换字符序列(字符串)
            @Override
            public CharSequence getCharContent(boolean ignoreEncodingErrors) {
                return code;
            }
        };

        // 获取系统Java编译器
        JavaCompiler compiler = ToolProvider.getSystemJavaCompiler();
        // 设置诊断收集器
        DiagnosticCollector<JavaFileObject> diagnostic = new DiagnosticCollector<>();
        // 获取Java文件管理器
        StandardJavaFileManager fileManager =
                compiler.getStandardFileManager(diagnostic, Locale.getDefault(), StandardCharsets.UTF_8);
        // 设置编译的操作选项["-d", "ojOutput":指定存放目录(相对路径)]
        Iterable<String> compilerOptions = Arrays.asList("-d", outDir, "-source", "21");
        // javaFileObject为设置编译单元
        Iterable<? extends JavaFileObject> compilerUnits = Collections.singletonList(javaFileObject);
        // getTask()获取待编译任务，call()执行编译
        JavaCompiler.CompilationTask task = compiler.getTask(
                null,
                fileManager,
                diagnostic,
                compilerOptions,
                null,
                compilerUnits);
        Boolean compileOk = task.call();
        // 关闭文件管理器资源
        try {
            fileManager.close();
        } catch (IOException e) {
            log.warn("Java管理器文件关闭失败：{}", e.toString());
        }
        // 编译失败
        if (!compileOk) {
            StringBuffer msg = new StringBuffer("\n--");
            diagnostic.getDiagnostics().forEach(dis -> {
                log.info("【{}】:【{}】", dis.getKind(), dis.getMessage(Locale.getDefault()));
                if (dis.getKind() == Diagnostic.Kind.ERROR)
                    msg.append(dis.getMessage(Locale.getDefault())).append("\n");
            });
            throw JudgeException.throwError(ExceptionConstant.COMPILER_ERROR + msg);
        }

    }

    /**
     * 加载类（字节码）文件并执行代码，比较原型结果和代码输出结果
     *
     * @param outDir      输出目录
     * @param className   类名
     * @param methodName  方法名
     * @param argsType    原型入惨类型
     * @param args        原型入惨
     * @param outcomeType 原型结果类型
     * @param outcome     原型结果
     * @return JudgedInfo
     */
    private JudgedInfo classLoadExecute(String outDir, String className, String methodName,
                                        Class<?>[] argsType, List<List<?>> args,
                                        Class<?> outcomeType, List<List<?>> outcome) {
        // 初始化使用的时间和内存，
        long useTimeNano = 1L; //纳秒 1ms = 1_000_000ns
        long useMemoryByte = 1L; //字节
        // 获取内存管理器
        MemoryMXBean memoryBean = ManagementFactory.getMemoryMXBean();
        // 初始化代码输出的结果
        List<List<?>> compiledResult = new ArrayList<>();
        // 初始化自定义安全类加载器
        MySecureClassLoader classLoader = null;
        try {
            // 获取编译后文件所在的目录路径的URL
            URL[] urls = new URL[]{Paths.get(outDir).toUri().toURL()};
            // 使用自定义类加载器
            classLoader = new MySecureClassLoader(urls);
            // 获取编译后的class，进行反射操作
            Class<?> clazz = classLoader.loadClass(className);
            Object instance = clazz.getDeclaredConstructor().newInstance();
            Method method = clazz.getDeclaredMethod(methodName, argsType);
            // 开始获取代码输出结果
            for (List<?> list : args) {
                // 入参列表转换为数组
                Object[] ias = list.toArray();
                // 入参格式化
                for (int i = 0; i < ias.length; i++) {
                    ias[i] = ParamsTypeFormat.argvFormat(ias[i], argsType[i]);
                }
                // 开始方法的计时与内存
                long beginT = System.nanoTime();
                MemoryUsage beginMUsage = memoryBean.getHeapMemoryUsage();
                // 方法发射获取输出结果
                Object irs = method.invoke(instance, ias);
                // 结束方法的计时与内存
                long endT = System.nanoTime();
                MemoryUsage endMUsage = memoryBean.getHeapMemoryUsage();
                useTimeNano += (endT - beginT);
                useMemoryByte += (endMUsage.getUsed() - beginMUsage.getUsed());
                // 格式化并添加输出结果，此时输出结果直接存储可能为地址（数组等类型），需格式化
                compiledResult.add(List.of(this.invokeResultFormat(irs)));
            }
        } catch (Exception e) {
            throw JudgeException.throwError(ExceptionConstant.COMPILER_ERROR + e.getMessage());
        } finally {
            try {
                if (classLoader != null) {
                    classLoader.close();
                }
                Files.deleteIfExists(Paths.get(outDir, className + ".class"));
            } catch (Exception e) {
                log.warn(e.toString());
            }
        }
        // 得到比较结果
        JudgedInfo judgedInfo = JudgeWorker.compare(outcomeType, outcome, compiledResult);
        // 结果不正确
        if (!judgedInfo.getFlag()) {
            return judgedInfo;
        }
        // 结果正确
        Integer useTimeMill = Integer.valueOf(Long.toString(useTimeNano / 1_000_000)); //ms
        Integer useMemoryKb = Integer.valueOf(Long.toString(useMemoryByte / 1024)); //kb
        return JudgedInfo.of(useTimeMill, useMemoryKb);
    }

    /**
     * 反射结果格式化，将数组类型的值格式化为集合类型 Array->List
     *
     * @param obj Object
     * @return Object
     */
    private Object invokeResultFormat(Object obj) {
        /*
        注意：
        1.区分基本类型、引用类型、数组、列表和其他类型，
        2.基本数据类型数组不能转为Object[]，其他引用类型数组可以转为Object[]
        */
        // 为空或者不是数组类型
        if (null == obj || !obj.getClass().isArray()) {
            return obj;
        }
        // 是数组类型，初始化返回结果
        List<Object> rs = new ArrayList<>();
        // 基本数据类型
        if (obj instanceof int[] o) {
            for (int i : o) {
                rs.add(i);
            }
        } else if (obj instanceof long[] o) {
            for (long l : o) {
                rs.add(l);
            }
        } else if (obj instanceof float[] o) {
            for (float f : o) {
                rs.add(f);
            }
        } else if (obj instanceof double[] o) {
            for (double d : o) {
                rs.add(d);
            }
        } else if (obj instanceof boolean[] o) {
            for (boolean b : o) {
                rs.add(b);
            }
        } else if (obj instanceof char[] o) {
            for (char c : o) {
                rs.add(c);
            }
        } else {
            rs.addAll(Arrays.asList((Object[]) obj));
        }
        // 返回结果
        return rs;
    }

    /**
     * 自定义安全类加载器
     *
     * @author ZSJ
     */
    public static class MySecureClassLoader extends URLClassLoader {

        /**
         * 权限
         */
        private static final PermissionCollection AllowedPermissions;

        /**
         * 目录路径
         */
        private final String[] pathDirs;

        /*
        static初始化
        */
        static {
            // 定义允许的权限（仅允许反射等基础操作）
            AllowedPermissions = new Permissions();
            AllowedPermissions.add(new RuntimePermission("accessDeclaredMembers"));
            AllowedPermissions.add(new ReflectPermission("suppressAccessChecks"));
        }

        public MySecureClassLoader(URL[] urls, ClassLoader parent) {
            super(urls, parent);
            // 将所有的url转为pathFirst（文件目录）
            pathDirs = new String[urls.length];
            for (int i = 0; i < urls.length; i++) {
                File file = new File(urls[i].getPath());
                pathDirs[i] = file.getPath();
                if (file.exists() && !file.delete()) {
                    log.warn("类加载器-初始化时，无法删除:【{}】 ", file.getPath());
                }
            }
        }

        public MySecureClassLoader(URL[] urls) {
            // ClassLoader.getSystemClassLoader() 获取系统类加载器
            this(urls, ClassLoader.getSystemClassLoader());
        }


        @Override
        protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
//        if (!name.startsWith("com.your.safe.package")) { // 限制包名
//            throw new ClassNotFoundException("未经授权 class: " + name);
//        }
            // loadClass(String name)调用loadClass(String name, resolve = false);
            return super.loadClass(name, resolve);
        }

        @Override
        protected Class<?> findClass(String name) throws ClassNotFoundException {
            // 1. 从路径加载字节码（示例路径），获取字节码数组
            byte[] classBytes = loadClassBytes(name);
            // 2. 字节码安全检查（使用ASM库）
            if (!isClassSafe(classBytes)) {
                log.error("类加载器-方法检测到恶意代码: {}", name);
                throw new SecurityException("检测到恶意代码: " + name);
            }

            // 3. 定义类时绑定自定义ProtectionDomain
            // 3.1 获取this类(SecureClassLoader)的CodeSource
            CodeSource codeSource = new CodeSource(
                    Objects.requireNonNull(getClassLocation(this.getClass())).getLocation(), (CodeSigner[]) null);
            // 3.2 设置该路径的权限控制
            ProtectionDomain domain = new ProtectionDomain(codeSource, AllowedPermissions);

            // 4. 在defineClass时，该类继承(绑定到)特定的(SecureClassLoader)CodeSource和权限
            return defineClass(name, classBytes, 0, classBytes.length, domain);
            // 虽然该类(MyClass.class)的物理路径是其他的(pools/)，
            // 但通过CodeSource，ProtectionDomain，defineClass()的绑定，
            // 它的安全上下文将继承自特定的(SecureClassLoader)的来源。
        }

        /**
         * 获取字节码数组
         *
         * @param name 类名
         * @return byte[]
         * @throws ClassNotFoundException 文件未找到
         */
        private byte[] loadClassBytes(String name) throws ClassNotFoundException {
            for (String pathDir : this.pathDirs) {
                String pathname = pathDir + "/" + name.replace(".", "/") + ".class";
                // 会抛出IOException
                try (InputStream is = this.getClass().getClassLoader().getResourceAsStream(pathname)) {
                    if (is != null) return is.readAllBytes();

                    // NullPointerException：当 path 为 null 时抛出。
                    // FileNotFoundException[继承IOException]：当 path 指向不存在的文件时抛出。
                    FileInputStream fis = new FileInputStream(pathname);
                    byte[] bytes = fis.readAllBytes();
                    fis.close();

                    return bytes;

                } catch (Exception e) {
                    log.warn("MySecureClassLoader的loadClassBytes方法，{}", e.toString());
                }
            }

            // pathFirsts遍历完没return则不存在
            throw new ClassNotFoundException(name);
        }

        /**
         * 检查字节码是否安全
         *
         * @param bytes 字节数组
         * @return boolean
         */
        private boolean isClassSafe(byte[] bytes) {
            // 使用ASM检查字节码（示例：禁止Runtime调用）
            ClassReader reader = new ClassReader(bytes);
            reader.accept(new ClassVisitor(Opcodes.ASM9) {
                @Override
                public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
                    return new MethodVisitor(Opcodes.ASM9) {
                        @Override
                        public void visitMethodInsn(int opcode, String owner, String methodName, String descriptor, boolean isInterface) {
                            if (owner.equals("java/lang/Runtime") /*&& methodName.equals("exec")*/) {
                                throw new SecurityException("禁止的调用: Runtime" /*Runtime.exec*/);
                            }
                            if (owner.equals("java/io/File")/* && methodName.equals("<init>")*/) {
                                throw new SecurityException("文件访问被禁止");
                            }
                            if (owner.startsWith("java/net/Socket")) {
                                throw new SecurityException("网络访问被禁止");
                            }
                        }
                    };
                }
            }, 0);

            return true;
        }

        /**
         * 获取代码源
         *
         * @param clazz 类
         * @return CodeSource
         */
        private CodeSource getClassLocation(Class<?> clazz) {
            ProtectionDomain domain = clazz.getProtectionDomain();
            return null != domain ? domain.getCodeSource() : null;
        }

    }


}

